import torch
from torchtext.datasets import AG_NEWS
from torchtext.data.utils import get_tokenizer
from collections import Counter
from torchtext.vocab import Vocab
from torch.utils.data import DataLoader
from torch import nn
from torch.optim import Optimizer as optimizer
import time
from torch.utils.data.dataset import random_split

"""
读取数据集
"""

# train_iter为RawTextIterableDataset类型
train_iter, test_iter = AG_NEWS(root='../data', split=('train', 'test'))
# print(type(train_iter))
# print(next(train_iter), next(test_iter))

"""
数据预处理
       1. 对每个word计数，并保存到dict里，类似于索引字典，使用collections.Counter实现，因为Vocab需要
       然后用预训练的模型转为词向量
"""

tokenizer = get_tokenizer(tokenizer='basic_english', language='en')
counter = Counter()
for (label, line) in train_iter:
    counter.update(tokenizer(line))
# print(counter)
vocab = Vocab(counter=counter, min_freq=1)
# vocab.load_vectors(vectors='fasttext.simple.300d')


# x = [vocab[token] for token in ['here', 'is', 'an', 'example']]
# print(x)


def text_to_num(text):
    return [vocab[token] for token in tokenizer(text)]


def label_to_num(y):
    return int(y) - 1


"""
生成数据批处理和迭代器
"""

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def collate_batch(batch):
    text_list, label_list, offsets = [], [], [0]  # offsets为偏移量（偏置）
    for (_label, _text) in batch:
        label_list.append(label_to_num(_label))
        processed_text = torch.tensor(data=text_to_num(_text), dtype=torch.int64)
        text_list.append(processed_text)
        print(len(processed_text))
        offsets.append(processed_text.size(0))
    label_list = torch.tensor(data=label_list, dtype=torch.int64)
    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
    text_list = torch.cat(tensors=text_list)  # cat连接张量tensor
    return label_list.to(device), text_list.to(device), offsets.to(device)  # 注意返回数据的顺序要与train模块加载的dataloader一致


# dataloader = DataLoader(dataset=train_iter, batch_size=32, shuffle=False, collate_fn=collate_batch)

"""
定义模型
"""


class TextClassificationModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, num_class):
        super(TextClassificationModel, self).__init__()
        self.embedding = nn.EmbeddingBag(num_embeddings=vocab_size, embedding_dim=embed_dim,
                                         sparse=True)  # sparse=True代表每次对该层求解梯度时, 只更新部分权重.
        self.fc = nn.Linear(embed_dim, num_class)  # 全连接层
        self.init_weights()

    def forward(self, text, offsets):
        embedded = self.embedding(text, offsets)
        return self.fc(embedded)

    def init_weights(self):
        initrange = 0.5
        # 权重初始化为平均分布
        self.embedding.weight.data.uniform_(-initrange, initrange)
        self.fc.weight.data.uniform_(-initrange, initrange)
        self.fc.bias.data.zero_()


"""
定义训练模型、评估模型方法
"""


def train(dataloader):
    model.train()
    total_acc, total_count = 0, 0
    log_interval = 500
    start_time = time.time()

    for idx, (label, text, offsets) in enumerate(dataloader):
        optimizer.zero_grad()
        predited_label = model(text, offsets)  # 输入text，offsets，获得预测值predited_label
        loss = criterion(predited_label, label)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)  # 梯度剪裁，防止梯度爆炸
        optimizer.step()  # 参数更新

        # 将该次的准确率加入到总的准确率中（这里计算的其实只是对的个数，最后还要除以所有样本总数）
        total_acc += (predited_label.argmax(1) == label).sum().item()
        total_count += label.size(0)
        if idx % log_interval == 0 and idx > 0:
            elapsed = time.time() - start_time
            print('| epoch {:3d} | {:5d}/{:5d} batches '
                  '| accuracy {:8.3f}'.format(epoch, idx, len(dataloader),
                                              total_acc / total_count))
            total_acc, total_count = 0, 0
            start_time = time.time()


def evaluate(dataloader):
    model.eval()
    total_acc, total_count = 0, 0
    with torch.no_grad():  # 验证阶段无需求解梯度
        for idx, (label, text, offsets) in enumerate(dataloader):
            predited_label = model(text, offsets)
            loss = criterion(predited_label, label)
            total_acc += (predited_label.argmax(1) == label).sum().item()
            total_count += label.size(0)
    return total_acc / total_count


"""
预测
"""


def predict(text, text_pipeline):
    with torch.no_grad():
        text = torch.tensor(text_pipeline(text))
        output = model(text, torch.tensor([0]))
        return output.argmax(1).item() + 1


"""
初始化实例
"""
# num_class = len(set([label for (label, text) in train_iter]))
num_class = 4
vocab_size = len(vocab)
emsize = 64
model = TextClassificationModel(vocab_size, emsize, num_class).to(device)

"""
划分数据集并运行
"""

if __name__ == '__main__':
    # Hyperparameters
    EPOCHS = 10  # epoch
    LR = 5  # learning rate
    BATCH_SIZE = 64  # batch size for training
    train_iter, test_iter = AG_NEWS(root='../data')
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=LR)
    # 调节学习率。lr = lr * gamma
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.1)
    total_accu = None

    # 切割数据集
    train_dataset = list(train_iter)
    test_dataset = list(test_iter)
    num_train = int(len(train_dataset) * 0.95)
    split_train_, split_valid_ = random_split(train_dataset, [num_train, len(train_dataset) - num_train])

    train_dataloader = DataLoader(split_train_, batch_size=BATCH_SIZE,
                                  shuffle=False, collate_fn=collate_batch)
    valid_dataloader = DataLoader(split_valid_, batch_size=BATCH_SIZE,
                                  shuffle=False, collate_fn=collate_batch)
    test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE,
                                 shuffle=False, collate_fn=collate_batch)

    for epoch in range(1, EPOCHS + 1):
        epoch_start_time = time.time()
        train(train_dataloader)
        accu_val = evaluate(valid_dataloader)
        if total_accu is not None and total_accu > accu_val:
            scheduler.step()
        else:
            total_accu = accu_val
        print('-' * 59)
        print('| end of epoch {:3d} | time: {:5.2f}s | '
              'valid accuracy {:8.3f} '.format(epoch,
                                               time.time() - epoch_start_time,
                                               accu_val))
        print('-' * 59)
    print('Checking the results of test dataset.')
    accu_test = evaluate(test_dataloader)
    print('test accuracy {:8.3f}'.format(accu_test))

    # 输出最后训练好的embedding weights矩阵
    print(model.state_dict()['embedding.weight'])

    ag_news_label = {1: "World",
                     2: "Sports",
                     3: "Business",
                     4: "Sci/Tec"}

    ex_text_str = "MEMPHIS, Tenn. – Four days ago, Jon Rahm was \
        enduring the season’s worst weather conditions on Sunday at The \
        Open on his way to a closing 75 at Royal Portrush, which \
        considering the wind and the rain was a respectable showing. \
        Thursday’s first round at the WGC-FedEx St. Jude Invitational \
        was another story. With temperatures in the mid-80s and hardly any \
        wind, the Spaniard was 13 strokes better in a flawless round. \
        Thanks to his best putting performance on the PGA Tour, Rahm \
        finished with an 8-under 62 for a three-stroke lead, which \
        was even more impressive considering he’d never played the \
        front nine at TPC Southwind."

    model = model.to("cpu")
    torch.save(model, 'model.pkl')
    print("This is a %s news" % ag_news_label[predict(ex_text_str, text_to_num)])
